package com.share.admin.web;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.share.admin.config.ChatGptConfig;
import com.share.admin.model.vo.OpenAiReq;
import com.share.common.common.Result;
import com.share.common.utils.HttpUtil;
import com.share.config.redis.RedisClient;
import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.service.OpenAiService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@RestController
@RequestMapping("/openai")
public class ChatGptController {
    @Autowired
    private OpenAiService openAiService;
    @Autowired
    private RedisClient redisClient;

    private static Map<String, String> header = new HashMap<>();

    static {
        header.put("Authorization", "Bearer " + ChatGptConfig.OPEN_AI_TOKEN);//中间有个token
        header.put("Content-Type", "application/json");
    }

    /**
     * 支持一个模型
     *
     * @return
     */
    @PostMapping("/h/ask")
    public Result askByHttp(@RequestBody OpenAiReq req) {
        long startTime = System.currentTimeMillis();
        log.info("askByJar-req={}", JSON.toJSONString(req));
        String url = "https://api.openai.com/v1/completions";
        JSONObject param = new JSONObject();
        param.put("prompt", req.getQuestion());//提问的内容
        param.put("max_tokens", 2048);
        param.put("model", "text-davinci-003");//模型

        try {
            String res = HttpUtil.post(url, param.toJSONString(), header, 60 * 1000);
            long endTime = System.currentTimeMillis();
            log.info("askByJar-costTimef秒={},question={},result={}", endTime - startTime, req.getQuestion(), JSON.toJSONString(res));

            //存入redis中
            String key = "chatgptj:" + startTime + "_" + req.getQuestion().trim();
            redisClient.set(key, JSON.toJSONString(res), Duration.ofDays(3l).toMinutes() * 60);

            return Result.success(res);
        } catch (Exception e) {
            log.error("请求异常", e);
            return Result.failure(e.getMessage());
        }
    }

    @PostMapping("/j/ask")
    public Result askByJar(@RequestBody OpenAiReq req) {
        long startTime = System.currentTimeMillis();
        log.info("askByJar-req={}", JSON.toJSONString(req));
        String model = req.getModel() == null ? ChatGptConfig.DEFAULT_MODEL : req.getModel();
        Double temperature = req.getTemperature() == null ? ChatGptConfig.DEFAULT_TEMPERATURE : req.getTemperature();
        CompletionRequest.CompletionRequestBuilder p = CompletionRequest.builder();
        CompletionRequest request = p.prompt(req.getQuestion()).model(model).maxTokens(2048).temperature(temperature).user(req.getUser()).build();
        List<CompletionChoice> choices = openAiService.createCompletion(request).getChoices();
        List<String> list = choices.stream().map(e -> e.getText()).collect(Collectors.toList());
        long endTime = System.currentTimeMillis();
        log.info("askByJar-costTimef秒={},question={},result={}", endTime - startTime, req.getQuestion(), JSON.toJSONString(list));

        //存入redis中
        String key = "chatgptj:" + startTime + "_" + req.getQuestion().trim();
        redisClient.set(key, JSON.toJSONString(list), Duration.ofDays(3l).toMinutes() * 60);

        return Result.success(list);
    }
}